#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
基于 OpenCV 的快速 OCR 解决方案
当 Tesseract 不可用时，提供快速的替代方案
"""

import cv2
import numpy as np
import re

class FastOpenCVOCR:
    """基于 OpenCV 的快速 OCR 识别器"""
    
    def __init__(self):
        self.digit_templates = self._create_digit_templates()
        self.letter_templates = self._create_letter_templates()
        print("✓ OpenCV OCR 已初始化")
    
    def _create_digit_templates(self):
        """创建数字模板"""
        templates = {}
        digits = "0123456789"
        
        for digit in digits:
            # 创建数字图像
            img = np.ones((50, 30), dtype=np.uint8) * 255
            cv2.putText(img, digit, (5, 35), cv2.FONT_HERSHEY_SIMPLEX, 1, 0, 2)
            templates[digit] = img
        
        return templates
    
    def _create_letter_templates(self):
        """创建字母模板"""
        templates = {}
        letters = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
        
        for letter in letters:
            # 创建字母图像
            img = np.ones((50, 30), dtype=np.uint8) * 255
            cv2.putText(img, letter, (5, 35), cv2.FONT_HERSHEY_SIMPLEX, 0.8, 0, 2)
            templates[letter] = img
        
        return templates
    
    def preprocess_image(self, roi):
        """预处理图像以提高识别效果"""
        try:
            # 转换为灰度图
            if len(roi.shape) == 3:
                gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
            else:
                gray = roi
            
            # 确保图像尺寸足够大
            height, width = gray.shape
            min_height, min_width = 50, 30
            
            if height < min_height or width < min_width:
                # 计算缩放比例，确保图像足够大
                scale = max(min_height / height, min_width / width)
                new_height = max(min_height, int(height * scale))
                new_width = max(min_width, int(width * scale))
                gray = cv2.resize(gray, (new_width, new_height))
            
            # 应用高斯模糊去噪
            blurred = cv2.GaussianBlur(gray, (3, 3), 0)
            
            # 自适应阈值处理
            thresh = cv2.adaptiveThreshold(
                blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
                cv2.THRESH_BINARY, 11, 2
            )
            
            # 形态学操作清理噪点
            kernel = np.ones((2, 2), np.uint8)
            cleaned = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel)
            
            return cleaned
            
        except Exception as e:
            print(f"图像预处理错误: {e}")
            return None
    
    def template_matching(self, image, templates, threshold=0.7):
        """模板匹配识别"""
        results = []
        
        # 检查图像尺寸
        img_height, img_width = image.shape[:2]
        
        for char, template in templates.items():
            # 检查模板尺寸
            template_height, template_width = template.shape[:2]
            
            # 确保图像尺寸大于等于模板尺寸
            if img_height < template_height or img_width < template_width:
                # 如果图像太小，跳过这个模板
                continue
            
            try:
                # 模板匹配
                result = cv2.matchTemplate(image, template, cv2.TM_CCOEFF_NORMED)
                _, max_val, _, max_loc = cv2.minMaxLoc(result)
                
                if max_val > threshold:
                    results.append((char, max_val, max_loc))
            except cv2.error as e:
                # 如果模板匹配失败，跳过这个模板
                continue
        
        # 按置信度排序
        results.sort(key=lambda x: x[1], reverse=True)
        
        return results
    
    def recognize_digits(self, processed_image):
        """识别数字"""
        results = self.template_matching(processed_image, self.digit_templates, threshold=0.6)
        if results:
            return results[0][0]  # 返回置信度最高的数字
        return ""
    
    def recognize_letters(self, processed_image):
        """识别字母"""
        results = self.template_matching(processed_image, self.letter_templates, threshold=0.6)
        if results:
            return results[0][0]  # 返回置信度最高的字母
        return ""
    
    def recognize_text(self, roi):
        """识别文本内容"""
        try:
            # 检查ROI是否有效
            if roi is None or roi.size == 0:
                return "图像无效"
            
            # 预处理图像
            processed_image = self.preprocess_image(roi)
            
            # 检查预处理后的图像
            if processed_image is None or processed_image.size == 0:
                return "预处理失败"
            
            # 尝试识别数字
            digit_result = self.recognize_digits(processed_image)
            if digit_result:
                return digit_result
            
            # 尝试识别字母
            letter_result = self.recognize_letters(processed_image)
            if letter_result:
                return letter_result
            
            # 如果都失败，尝试轮廓检测
            return self._contour_based_recognition(processed_image)
            
        except Exception as e:
            return f"识别错误: {str(e)}"
    
    def _contour_based_recognition(self, processed_image):
        """基于轮廓的识别"""
        try:
            # 查找轮廓
            contours, _ = cv2.findContours(
                processed_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
            )
            
            if not contours:
                return "未识别"
            
            # 分析轮廓特征
            text_parts = []
            for contour in contours:
                # 计算轮廓面积
                area = cv2.contourArea(contour)
                if area < 50:  # 过滤太小的轮廓
                    continue
                
                # 获取边界矩形
                x, y, w, h = cv2.boundingRect(contour)
                
                # 根据宽高比判断可能的字符类型
                aspect_ratio = w / h if h > 0 else 0
                
                if 0.2 < aspect_ratio < 2.0:  # 合理的字符宽高比
                    # 提取ROI进行进一步分析
                    roi = processed_image[y:y+h, x:x+w]
                    
                    # 尝试识别
                    digit = self.recognize_digits(roi)
                    if digit:
                        text_parts.append(digit)
                        continue
                    
                    letter = self.recognize_letters(roi)
                    if letter:
                        text_parts.append(letter)
                        continue
            
            if text_parts:
                return ''.join(text_parts)
            
            return "未识别"
            
        except Exception as e:
            return f"轮廓识别错误: {str(e)}"
    
    def recognize_number(self, roi):
        """兼容接口，识别数字"""
        return self.recognize_text(roi)
    
    def benchmark(self, test_image, iterations=10):
        """性能测试"""
        import time
        
        start_time = time.time()
        for _ in range(iterations):
            self.recognize_text(test_image)
        elapsed = time.time() - start_time
        
        avg_time = elapsed / iterations
        print(f"OpenCV OCR 平均识别时间: {avg_time:.4f} 秒/次")
        return avg_time

# 使用示例
if __name__ == "__main__":
    # 创建快速 OCR 实例
    ocr = FastOpenCVOCR()
    
    # 创建测试图像
    test_img = np.ones((100, 200, 3), dtype=np.uint8) * 255
    cv2.putText(test_img, "123", (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2)
    cv2.putText(test_img, "ABC", (50, 80), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2)
    
    # 测试识别
    result = ocr.recognize_text(test_img)
    print(f"测试识别结果: {result}")
    
    # 性能测试
    ocr.benchmark(test_img, iterations=5)
